package svsim

import java.io.{BufferedReader, BufferedWriter, File, FileWriter, InputStreamReader, PrintWriter}
import java.nio.file.Paths
import java.lang.ProcessBuilder.Redirect
import scala.annotation.meta.param

case class ModuleInfo(
  name:  String,
  ports: Seq[ModuleInfo.Port]) {
  private[svsim] val instanceName = "dut"
}
object ModuleInfo {
  case class Port(
    name:       String,
    isSettable: Boolean = false,
    isGettable: Boolean = false) {
    assert(name.matches("^[a-zA-Z0-9\\-_]*$"))
  }
}

object Workspace {
  val testbenchModuleName: String = "svsimTestbench"
}
final class Workspace(
  path: String,
  /** The prefix for the working directory used when invoking `compile`
    */
  val workingDirectoryPrefix: String = "workdir") {

  val absolutePath =
    if (path.startsWith("/"))
      path
    else
      s"${System.getProperty("user.dir")}/$path"

  /** A directory where the user can store additional artifacts which are relevant to the primary sources (for instance, artifacts related to the generation of primary sources). These artifacts have no impact on the simulation, but it may be useful to group them with the other files generated by svsim for debugging purposes.
    */
  val supportArtifactsPath = s"$absolutePath/support-artifacts"

  /** The directory containing user-provided source files used to compile the simulation when `compile` is called.
    */
  val primarySourcesPath = s"$absolutePath/primary-sources"

  /** The directory containing code generated when calling `generateAdditionalSources`
    */
  val generatedSourcesPath = s"$absolutePath/generated-sources"

  private var _moduleInfo: Option[ModuleInfo] = None

  def reset() = {
    _moduleInfo = None

    val rm = Runtime.getRuntime().exec(Array("rm", "-rf", absolutePath)).waitFor()
    val pathsToCreate = Seq(
      supportArtifactsPath,
      primarySourcesPath,
      generatedSourcesPath
    )
    pathsToCreate.map(new File(_)).foreach(_.mkdirs())
  }

  private def copyResource(klass: Class[_], name: String, targetDirectory: String) = {
    val inputStream = klass.getResourceAsStream(name)
    if (inputStream == null) throw new java.io.FileNotFoundException(name)
    val file = new java.io.File(targetDirectory, name.split("/").last)
    val outputStream = new java.io.FileOutputStream(file)
    Iterator
      .continually(inputStream.read)
      .takeWhile(_ != -1)
      .foreach(outputStream.write)
    inputStream.close()
    outputStream.close()
  }

  /** A helper method which copies the specified resource into the primary sources directory.
    */
  def addPrimarySourceFromResource(klass: Class[_], name: String) = {
    copyResource(klass, name, primarySourcesPath)
  }

  /** `svsim` elaboration simply stores the provided `ModuleInfo` for use by the `compile` method. The idea is that packages that actually do elaboration (like Chisel) will add an overload of this method in an implicit class that then calls this method with the appropriate `ModuleInfo`.
    */
  def elaborate(moduleInfo: ModuleInfo) = {
    assert(_moduleInfo.isEmpty)
    _moduleInfo = Some(moduleInfo)
  }

  /** Generate additional sources necessary for simulating the module.
    */
  //format: off
  final def generateAdditionalSources() = {
    val dut = _moduleInfo.get
    val ports = dut.ports.zipWithIndex

    val systemVerilogTestbenchWriter = new LineWriter(s"$generatedSourcesPath/testbench.sv")
    try {
      val l = systemVerilogTestbenchWriter

      l("module ", Workspace.testbenchModuleName, ";")
      for ((port, index) <- ports) {
        if (port.isSettable) {
      l("  reg  [$bits(", dut.instanceName, ".",  port.name, ")-1:0] ", port.name, ";")
        } else {
      l("  wire [$bits(", dut.instanceName, ".",  port.name, ")-1:0] ", port.name, ";")
        }
      }
      l()
      l(dut.name, " ", dut.instanceName, " (")
      for ((port, index) <- ports) {
      l("    .", port.name, "(", port.name, ")", if (index != ports.length - 1) "," else "")
      }
      l(");")
      l()
      for ((port, index) <- ports) {
      l("  // Port ", index.toHexString, ": ", port.name)
      l("  export \"DPI-C\" function getBitWidth_", port.name, ";")
      l("  function void getBitWidth_", port.name, ";")
      l("    output int value;")
      l("    value", " = $bits(", dut.instanceName, ".",  port.name, ");")
      l("  endfunction")
        if (port.isSettable) {
      l("  export \"DPI-C\" function setBits_", port.name, ";")
      l("  function void setBits_", port.name, ";")
      l("    input bit [$bits(", dut.instanceName, ".",  port.name, ")-1:0] value_", port.name, ";")
      l("    ", port.name, " = value_", port.name, ";")
      l("  endfunction")
        }
        if (port.isGettable) {
      l("  export \"DPI-C\" function getBits_", port.name, ";")
      l("  function void getBits_", port.name, ";")
      l("    output bit [$bits(", dut.instanceName, ".",  port.name, ")-1:0] value_", port.name, ";")
      l("    value_", port.name, " = ", port.name, ";")
      l("  endfunction")
        }
      l()
      }
      l("  // Simulation")
      l("  import \"DPI-C\" context task simulation_body();")
      l("  initial begin")
      l("    simulation_body();")
      l("  end")
      l("  `ifdef ", Backend.HarnessCompilationFlags.supportsDelayInPublicFunctions)
      l("  export \"DPI-C\" task run_simulation;")
      l("  task run_simulation;")
      l("    input int timesteps;")
      l("    #timesteps;")
      l("  endtask")
      l("  `else")
      l("  import \"DPI-C\" function void run_simulation(int timesteps);")
      l("  `endif")
      l()

      l("  // Tracing")
      l("  export \"DPI-C\" function simulation_initializeTrace;")
      l("  function void simulation_initializeTrace;")
      l("    input string traceFilePath;")
      l("    `ifdef SVSIM_ENABLE_VCD_TRACING_SUPPORT")
      l("      $dumpfile({traceFilePath,\".vcd\"});")
      l("      $dumpvars(0, ", dut.instanceName,");")
      l("    `endif")
      l("    `ifdef SVSIM_ENABLE_VPD_TRACING_SUPPORT")
      l("      $vcdplusfile({traceFilePath,\".vpd\"});")
      l("      $dumpvars(0, ", dut.instanceName,");")
      l("      $vcdpluson(0, ", dut.instanceName,");")
      l("    `endif")
      l("    `ifdef SVSIM_ENABLE_FSDB_TRACING_SUPPORT")
      l("      $fsdbDumpfile({traceFilePath,\".fsdb\"});")
      l("      $fsdbDumpvars(0, ", dut.instanceName,");")
      l("    `endif")
      l("  endfunction")
      l("  export \"DPI-C\" function simulation_enableTrace;")
      l("  function void simulation_enableTrace;")
      l("    `ifdef SVSIM_ENABLE_VCD_TRACING_SUPPORT")
      l("    $dumpon;")
      l("    `elsif SVSIM_ENABLE_VPD_TRACING_SUPPORT")
      l("    $dumpon;")
      l("    `endif")
      l("    `ifdef SVSIM_ENABLE_FSDB_TRACING_SUPPORT")
      l("    $fsdbDumpon;")
      l("    `endif")
      l("  endfunction")
      l("  export \"DPI-C\" function simulation_disableTrace;")
      l("  function void simulation_disableTrace;")
      l("    `ifdef SVSIM_ENABLE_VCD_TRACING_SUPPORT")
      l("    $dumpoff;")
      l("    `elsif SVSIM_ENABLE_VPD_TRACING_SUPPORT")
      l("    $dumpoff;")
      l("    `endif")
      l("    `ifdef SVSIM_ENABLE_FSDB_TRACING_SUPPORT")
      l("    $fsdbDumpoff;")
      l("    `endif")
      l("  endfunction")
      l()
      l("endmodule")
    } finally {
      // `BufferedWriter` closes the underlying `FileWriter` when closed.
      systemVerilogTestbenchWriter.close()
    }

    val cDPIBridgeWriter = new LineWriter(s"$generatedSourcesPath/c-dpi-bridge.cpp")
    try {
      val l = cDPIBridgeWriter
      l("#include <stdint.h>")
      l()      
      l("#ifdef SVSIM_ENABLE_VERILATOR_SUPPORT")
      l("#include \"verilated-sources/VsvsimTestbench__Dpi.h\"")
      l("#endif")
      l("#ifdef SVSIM_ENABLE_VCS_SUPPORT")
      l("#include \"vc_hdrs.h\"")
      l("#endif")
      l()
      l("extern \"C\" {")
      l()
      l("int port_getter(int id, int *bitWidth, void (**getter)(uint8_t*)) {")
      l("  switch (id) {")
      for ((port, index) <- ports.filter(_._1.isGettable)) {
      l("    case ", index.toString(), ": // ", port.name)
      l("      getBitWidth_", port.name, "(bitWidth);")
      l("      *getter = (void(*)(uint8_t*))getBits_", port.name, ";")
      l("      return 0;")
      }
      l("    default:")
      l("      return -1;")
      l("  }")
      l("}")
      l()
      l("int port_setter(int id, int *bitWidth, void (**setter)(const uint8_t*)) {")
      l("  switch (id) {")
      for ((port, index) <- ports.filter(_._1.isSettable)) {
      l("    case ", index.toString(), ": // ", port.name)
      l("      getBitWidth_", port.name, "(bitWidth);")
      l("      *setter = (void(*)(const uint8_t*))setBits_", port.name, ";")
      l("      return 0;")
      }
      l("    default:")
      l("      return -1;")
      l("  }")
      l("}")
      l()
      l("} // extern \"C\"")
      l()

    } finally {
      cDPIBridgeWriter.close()
    }
  
    copyResource(this.getClass, "/simulation-driver.cpp", generatedSourcesPath)
  }
  //format: on

  /** Compiles the simulation using the specified backend.
    *
    * @param outputTag A string which will be used to tag the output directory. This enables compiling and simulating the same workspace with multiple backends.
    */
  def compile[T <: Backend](
    backend:                          T
  )(workingDirectoryTag:              String,
    commonSettings:                   CommonCompilationSettings,
    backendSpecificSettings:          backend.CompilationSettings,
    customSimulationWorkingDirectory: Option[String],
    verbose:                          Boolean
  ): Simulation = {
    val moduleInfo = _moduleInfo.get
    val workingDirectoryPath = s"$absolutePath/$workingDirectoryPrefix-$workingDirectoryTag"
    val workingDirectory = new File(workingDirectoryPath)
    workingDirectory.mkdir()

    val parameters = backend.generateParameters(
      outputBinaryName = "simulation",
      topModuleName = Workspace.testbenchModuleName,
      additionalHeaderPaths = Seq(workingDirectoryPath),
      commonSettings = commonSettings,
      backendSpecificSettings = backendSpecificSettings
    )
    val sourceFiles = Seq(primarySourcesPath, generatedSourcesPath)
      .map(new File(_))
      .flatMap(_.listFiles())
      .map { file => workingDirectory.toPath().relativize(file.toPath()).toString() }

    val simulationEnvironment = Seq(
      "SVSIM_SIMULATION_LOG" -> s"$workingDirectoryPath/simulation-log.txt",
      // The simulation driver appends the appropriate extension to the file path
      "SVSIM_SIMULATION_TRACE" -> s"$workingDirectoryPath/trace"
    ) ++ parameters.simulationInvocation.environment

    // Emit Makefile for debugging (will be emitted even if compile fails)
    val makefileWriter = new LineWriter(s"$workingDirectoryPath/Makefile")
    try {
      val l = makefileWriter
      l("# This Makefile enables lightweight debugging of `svsim` tests. ")
      l("# To rebuild the simulation run `make simulation` in this directory.")
      l("# To replay the simulation run `make replay` in this directory.")
      l(
        "# Changes to `generated-sources` and `primary-sources` will be picked up when running `make replay` and `make simulation`. This is useful for debugging issues. You can also freely add, remove or change any of the arguments to the backend or simulation in the targets below."
      )
      l()
      l(".PHONY: clean simulation replay")
      l()
      //format: off
      // For this debug flow, we rebuild the simulation from scratch every time, to avoid issues if the simulation was originally compiled in a different environment, like using SiFive's `wake`.
      l("clean:")
	    l("\tls . | grep -v Makefile | grep -v execution-script.txt | xargs rm -rf")
      l()
      l("simulation: clean")
      l("\t$(compilerEnvironment) \\")
	    l("\t", parameters.compilerPath, " \\")
      for (argument <- parameters.compilerInvocation.arguments) {
        val sanitizedArugment = argument
          .replace("$", "$$")
          .replace("'", "'\\''")
          .replace(workingDirectoryPath, "$(shell pwd)")
      l("\t\t'", sanitizedArugment, "' \\")
      }
      l("\t\t$(sourcefiles)")
      l()
      l("replay: simulation")
      val executionScriptPath = "$(shell pwd)/execution-script.txt"
      customSimulationWorkingDirectory match {
        case None =>
        case Some(value) => {
          // Calculate relative path, to avoid being broken by wake's directory shenanigans
          val relativePath = workingDirectory.toPath().relativize(Paths.get(value)).toString()
          l("\tcd ", relativePath, " && \\")
        }
      }
      l("\tcat ", executionScriptPath, " | { grep '^#' || true; } && \\")
      l("\tcat ", executionScriptPath, " | sed -n 's/^[0-9]*> \\(.*\\)/\\1/p' | \\")
      l("\t\t$(simulationEnvironment) $(shell pwd)/simulation \\")
      for (argument <- parameters.simulationInvocation.arguments) {
      l("\t\t\t'", argument.replace("$", "$$"), "' \\")
      }
      l()
      l("sourcefiles = \\")
      for ((sourceFile, index) <- sourceFiles.zipWithIndex) {
      l("\t'", sourceFile, "'", if (index != sourceFiles.length - 1) " \\" else "")
      }
      l()
      l("compilerEnvironment = \\")
      for (((name, value), index) <- parameters.compilerInvocation.environment.zipWithIndex) {
      l("\t", name, "=", value, if (index != parameters.compilerInvocation.environment.length - 1) " \\" else "")
      }
      l()
      l("simulationEnvironment = \\")
      for (((name, value), index) <- simulationEnvironment.zipWithIndex) {
        val sanitizedValue = value.replace(workingDirectoryPath, "$(shell pwd)")
      l("\t", name, "=", sanitizedValue, if (index != simulationEnvironment.length - 1) " \\" else "")
      }
      l()
      //format: on
    } finally {
      makefileWriter.close()
    }

    /**
      * Use the generated Makefile to compile the simulation, since this exercises the Makefile codepath and makes it less likely that we will break `make replay`.
      */
    val processBuilder = new ProcessBuilder("make", "-C", workingDirectoryPath, "simulation")
    processBuilder.redirectErrorStream(true)
    val process = processBuilder.start()
    @scala.annotation.nowarn(
      "msg=Use `scala.jdk.CollectionConverters` instead"
    )
    def readLogLines() = {
      val sourceLocationRegex = "[\\./]*generated-sources/".r
      import scala.collection.JavaConverters._
      new BufferedReader(new InputStreamReader(process.getInputStream()))
        .lines()
        .map(sourceLocationRegex.replaceFirstIn(_, ""))
        .map { line =>
          if (verbose) {
            println(line)
          }
          line
        }
        .iterator()
        .asScala
        .toSeq
    }
    val compilationLogLines = readLogLines()
    process.waitFor()
    val compilationLogWriter = new PrintWriter(
      new BufferedWriter(
        new FileWriter(new File(s"$workingDirectoryPath/compilation-log.txt"))
      )
    )
    compilationLogLines.foreach(compilationLogWriter.println)
    compilationLogWriter.close()
    if (process.exitValue() != 0) {
      throw new Exception(compilationLogLines.mkString("\n"))
    }

    new Simulation(
      executableName = "simulation",
      settings = Simulation.Settings(
        customWorkingDirectory = customSimulationWorkingDirectory,
        arguments = parameters.simulationInvocation.arguments,
        environment = simulationEnvironment.toMap
      ),
      workingDirectoryPath = workingDirectoryPath,
      moduleInfo = moduleInfo
    )
  }

}

/** A micro-DSL for writing files.
  */
private class LineWriter(path: String) {
  private val wrapped = new BufferedWriter(new FileWriter(path, false))
  def apply(components: String*) = {
    components.foreach(wrapped.write)
    wrapped.newLine()
  }
  def close() = wrapped.close()
}
